import os
import logging
from typing import Dict, List, Optional, Union

from models.ollama_wrapper import OllamaWrapper
from config.settings import OUTPUT_DIR

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


class DocGenerator:
    """
    基于代码分析结果生成文档
    """
    def __init__(self, analyzer, graph_db):
        self.analyzer = analyzer
        self.graph_db = graph_db
        self.llm = OllamaWrapper()
        self._ensure_output_dir()
        
    def _ensure_output_dir(self) -> None:
        """确保输出目录存在"""
        if not os.path.exists(OUTPUT_DIR):
            os.makedirs(OUTPUT_DIR)
            logger.info(f"创建输出目录: {OUTPUT_DIR}")

    def generate_file_doc(self, file_path: str) -> str:
        """
        为指定文件生成文档
        """
        logger.info(f"为文件生成文档: {file_path}")
        
        try:
            # 获取文件上下文信息
            file_context = self.analyzer.get_context_for_file(file_path)
            if "error" in file_context:
                return f"Error: {file_context['error']}"
                
            # 准备提示词
            prompt = self._prepare_file_doc_prompt(file_path, file_context)
            
            # 使用LLM生成文档
            markdown_doc = self.llm.generate(prompt)
            
            # 保存生成的文档
            self._save_doc(file_path, markdown_doc)
            
            return markdown_doc
            
        except Exception as e:
            logger.error(f"生成文件文档时出错: {str(e)}")
            return f"文档生成失败: {str(e)}"
            
    def generate_file_doc_observable(self, file_path: str) -> str:
        """
        为指定文件生成文档，同时可观察生成过程
        """
        logger.info(f"为文件生成可观察文档: {file_path}")
        
        try:
            # 获取文件上下文信息
            file_context = self.analyzer.get_context_for_file(file_path)
            if "error" in file_context:
                return f"Error: {file_context['error']}"
                
            # 准备提示词
            prompt = self._prepare_file_doc_prompt(file_path, file_context)
            print("\n==== 提示词 ====")
            print(prompt[:500] + "..." if len(prompt) > 500 else prompt)
            print("\n==== 开始生成 (实时输出) ====")
            
            # 使用流式生成
            full_response = ""
            for chunk in self.llm.generate_stream(prompt):
                print(chunk, end="", flush=True)
                full_response += chunk
            
            print("\n==== 生成完成 ====\n")
            
            # 保存生成的文档
            self._save_doc(file_path, full_response)
            
            return full_response
            
        except Exception as e:
            logger.error(f"生成可观察文档时出错: {str(e)}")
            return f"文档生成失败: {str(e)}"

    def generate_function_doc(self, file_path: str, function_name: str) -> str:
        """
        为指定函数生成文档
        """
        logger.info(f"为函数生成文档: {file_path}:{function_name}")
        
        try:
            # 获取函数上下文信息
            function_context = self.graph_db.get_function_context(file_path, function_name)
            if "error" in function_context:
                return f"Error: {function_context['error']}"
                
            # 获取函数所在文件的内容
            file_content = self.analyzer.get_file_content(file_path)
                
            # 准备提示词
            prompt = self._prepare_function_doc_prompt(file_path, function_name, function_context, file_content)
            
            # 使用LLM生成文档
            markdown_doc = self.llm.generate(prompt)
            
            # 保存生成的文档
            output_file = f"{file_path.replace('/', '_')}__{function_name}.md"
            self._save_doc(output_file, markdown_doc)
            
            return markdown_doc
            
        except Exception as e:
            logger.error(f"生成函数文档时出错: {str(e)}")
            return f"文档生成失败: {str(e)}"
            
    def generate_function_doc_observable(self, file_path: str, function_name: str) -> str:
        """
        为指定函数生成文档，同时可观察生成过程
        """
        logger.info(f"为函数生成可观察文档: {file_path}:{function_name}")
        
        try:
            # 获取函数上下文信息
            function_context = self.graph_db.get_function_context(file_path, function_name)
            if "error" in function_context:
                return f"Error: {function_context['error']}"
                
            # 获取函数所在文件的内容
            file_content = self.analyzer.get_file_content(file_path)
                
            # 准备提示词
            prompt = self._prepare_function_doc_prompt(file_path, function_name, function_context, file_content)
            print("\n==== 提示词 ====")
            print(prompt[:500] + "..." if len(prompt) > 500 else prompt)
            print("\n==== 开始生成 (实时输出) ====")
            
            # 使用流式生成
            full_response = ""
            for chunk in self.llm.generate_stream(prompt):
                print(chunk, end="", flush=True)
                full_response += chunk
            
            print("\n==== 生成完成 ====\n")
            
            # 保存生成的文档
            output_file = f"{file_path.replace('/', '_')}__{function_name}.md"
            self._save_doc(output_file, full_response)
            
            return full_response
            
        except Exception as e:
            logger.error(f"生成可观察函数文档时出错: {str(e)}")
            return f"文档生成失败: {str(e)}"

    def generate_project_overview(self) -> str:
        """
        生成项目概述文档
        """
        logger.info("生成项目概述文档")
        
        try:
            # 统计项目概览
            total_files = len(self.analyzer.files)
            file_types = {}
            for file in self.analyzer.files:
                ext = os.path.splitext(file)[1]
                if ext in file_types:
                    file_types[ext] += 1
                else:
                    file_types[ext] = 1
            
            # 准备提示词
            prompt = f"""
            基于以下项目信息生成一份全面的项目概述文档：
            
            项目统计信息：
            - 总文件数: {total_files}
            - 文件类型分布: {file_types}
            
            重要文件列表:
            {self._get_important_files()}
            
            主要模块和功能:
            {self._get_key_modules()}
            
            请生成一份Markdown格式的项目概述文档，包括：
            1. 项目架构概述
            2. 核心模块说明
            3. 关键功能与流程
            4. 技术栈分析
            5. 组件间依赖关系
            
            文档应该简洁明了，重点突出项目的整体结构和主要功能。
            """
            
            # 使用LLM生成文档
            markdown_doc = self.llm.generate(prompt)
            
            # 保存生成的文档
            self._save_doc("project_overview.md", markdown_doc)
            
            return markdown_doc
            
        except Exception as e:
            logger.error(f"生成项目概述时出错: {str(e)}")
            return f"项目概述生成失败: {str(e)}"
            
    def generate_project_overview_observable(self) -> str:
        """
        生成项目概述文档，同时可观察生成过程
        """
        logger.info("生成可观察项目概述文档")
        
        try:
            # 统计项目概览
            total_files = len(self.analyzer.files)
            file_types = {}
            for file in self.analyzer.files:
                ext = os.path.splitext(file)[1]
                if ext in file_types:
                    file_types[ext] += 1
                else:
                    file_types[ext] = 1
            
            # 准备提示词
            prompt = f"""
            基于以下项目信息生成一份全面的项目概述文档：
            
            项目统计信息：
            - 总文件数: {total_files}
            - 文件类型分布: {file_types}
            
            重要文件列表:
            {self._get_important_files()}
            
            主要模块和功能:
            {self._get_key_modules()}
            
            请生成一份Markdown格式的项目概述文档，包括：
            1. 项目架构概述
            2. 核心模块说明
            3. 关键功能与流程
            4. 技术栈分析
            5. 组件间依赖关系
            
            文档应该简洁明了，重点突出项目的整体结构和主要功能。
            """
            
            print("\n==== 提示词 ====")
            print(prompt[:500] + "..." if len(prompt) > 500 else prompt)
            print("\n==== 开始生成 (实时输出) ====")
            
            # 使用流式生成
            full_response = ""
            for chunk in self.llm.generate_stream(prompt):
                print(chunk, end="", flush=True)
                full_response += chunk
            
            print("\n==== 生成完成 ====\n")
            
            # 保存生成的文档
            self._save_doc("project_overview.md", full_response)
            
            return full_response
            
        except Exception as e:
            logger.error(f"生成可观察项目概述时出错: {str(e)}")
            return f"项目概述生成失败: {str(e)}"

    def _prepare_file_doc_prompt(self, file_path: str, context: Dict) -> str:
        """准备文件文档生成的提示词"""
        file_content = context.get("content", "")
        imports = context.get("imports", [])
        functions = context.get("functions", [])
        calls = context.get("calls", {})
        
        import_info = "\n".join([f"- {imp.get('module')} " + 
                                (f"as {imp.get('alias')}" if imp.get('alias') else "") 
                                for imp in imports])
        
        function_info = ""
        for func in functions:
            func_name = func.get("name", "")
            params_str = ", ".join(func.get("params", []))
            return_type = func.get("return_type", "")
            docstring = func.get("docstring", "")
            
            function_info += f"\n- {func_name}({params_str}){return_type}\n"
            if docstring:
                function_info += f"  Docstring: {docstring}\n"
                
            # 添加调用信息
            if func_name in calls:
                function_info += "  调用关系:\n"
                for call in calls[func_name]:
                    if call.get("type") == "local":
                        function_info += f"  - 调用本地函数: {call.get('name')} (在 {call.get('file')})\n"
                    else:
                        function_info += f"  - 调用外部函数: {call.get('module')}.{call.get('method')}\n"
        
        prompt = f"""
        请为以下代码文件生成一份详细的技术文档。

        文件路径: {file_path}
        
        导入关系:
        {import_info if import_info else "无导入"}
        
        函数概览:
        {function_info if function_info else "无函数定义"}
        
        文件内容:
        {file_content[:4000] if len(file_content) > 4000 else file_content}
        {"..." if len(file_content) > 4000 else ""}
                
        请生成一份Markdown格式的文档，包括：
        1. 文件概述：简要介绍该文件的主要功能和用途
        2. 依赖关系：该文件导入了哪些模块，如何使用它们
        3. 核心功能：详细说明各核心函数的功能和用途
        4. 内部工作原理：分析代码的工作流程和实现细节
        5. 使用示例：如何使用该文件中的功能（请根据代码推断）
        
        文档应该简洁清晰，注重功能解释而不是代码复述。
        """
        
        return prompt

    def _prepare_function_doc_prompt(self, file_path: str, function_name: str, 
                                    context: Dict, file_content: str) -> str:
        """准备函数文档生成的提示词"""
        params = context.get("params", [])
        return_type = context.get("return_type", "")
        docstring = context.get("docstring", "")
        calls = context.get("calls", [])
        called_by = context.get("called_by", [])
        
        # 提取函数代码（简单实现，实际应使用AST更准确）
        import re
        pattern = fr"def\s+{function_name}\s*$[^)]*$\s*(?:->\s*[^:]+\s*)?:"
        match = re.search(pattern, file_content)
        function_code = ""
        if match:
            start = match.start()
            # 简单提取，不处理嵌套函数
            level = 0
            lines = file_content[start:].split("\n")
            function_lines = []
            
            for i, line in enumerate(lines):
                if i == 0:
                    function_lines.append(line)
                    level = len(line) - len(line.lstrip())
                    continue
                    
                current_level = len(line) - len(line.lstrip())
                if line.strip() and current_level <= level and i > 1:
                    break
                function_lines.append(line)
                
            function_code = "\n".join(function_lines)
        
        calls_info = ""
        for call in calls:
            if call.get("external", False):
                calls_info += f"- 调用外部模块: {call.get('module')}.{call.get('name')}\n"
            else:
                calls_info += f"- 调用本地函数: {call.get('name')} (在 {call.get('file')})\n"
        
        called_by_info = ""
        for caller in called_by:
            called_by_info += f"- 被函数 {caller.get('name')} 调用 (在 {caller.get('file')})\n"
        
        prompt = f"""
        请为以下函数生成一份详细的技术文档。

        文件路径: {file_path}
        函数名: {function_name}
        
        函数签名: 
        {function_name}({", ".join(params)}){return_type}
        
        函数注释:
        {docstring if docstring else "无文档注释"}
        
        函数调用关系:
        {calls_info if calls_info else "不调用其他函数"}
        
        函数被调用情况:
        {called_by_info if called_by_info else "未被其他函数调用"}
        
        函数代码:
        python
        {function_code}
                
        请生成一份Markdown格式的函数文档，包括：
        1. 函数概述：简要介绍该函数的主要功能和用途
        2. 参数说明：详细说明各参数的类型、用途和限制
        3. 返回值说明：说明函数返回值的类型和含义
        4. 工作原理：分析函数的内部实现逻辑和算法
        5. 示例用法：提供调用该函数的示例代码
        6. 注意事项：使用该函数时需要注意的事项或潜在问题
        
        文档应该技术准确，易于理解，并考虑到函数与其他组件的交互关系。
        """
        
        return prompt

    def _get_important_files(self) -> str:
        """识别项目中的重要文件"""
        important_files = []
        
        # 启发式方法识别重要文件
        for file_path in self.analyzer.files:
            # 假设名称包含main、app、core、server等的文件更重要
            lower_path = file_path.lower()
            if any(keyword in lower_path for keyword in ['main', 'app', 'core', 'server', 'api']):
                important_files.append(file_path)
        
        # 按调用频率排序
        file_call_count = {}
        for rel in self.graph_db.call_relations:
            source = rel["source"].split(":")[0]
            if source in file_call_count:
                file_call_count[source] += 1
            else:
                file_call_count[source] = 1
        
        sorted_files = sorted(file_call_count.items(), key=lambda x: x[1], reverse=True)
        for file, count in sorted_files[:5]:  # 取前5个最活跃文件
            if file not in important_files:
                important_files.append(file)
        
        return "\n".join([f"- {file}" for file in important_files[:10]])  # 最多返回10个

    def _get_key_modules(self) -> str:
        """识别项目中的关键模块"""
        # 根据目录结构识别模块
        modules = set()
        for file_path in self.analyzer.files:
            parts = file_path.split('/')
            if len(parts) > 1:
                modules.add(parts[0])
        
        module_info = ""
        for module in modules:
            module_files = [f for f in self.analyzer.files if f.startswith(f"{module}/")]
            if len(module_files) > 0:
                module_info += f"- {module}/: 包含 {len(module_files)} 个文件\n"
                
                # 添加一些代表性文件
                sample_files = module_files[:3]  # 最多列出3个
                for sample in sample_files:
                    module_info += f"  - {sample}\n"
        
        return module_info

    def _save_doc(self, file_path: str, content: str) -> None:
        """保存生成的文档"""
        # 处理文件名，确保安全
        safe_filename = file_path.replace('/', '_').replace('\\', '_')
        output_path = os.path.join(OUTPUT_DIR, safe_filename)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(content)
            
        logger.info(f"已保存文档到: {output_path}")

    